import os
import shutil
from argparse import ArgumentParser

import cv2
import torch
import numpy as np
from scipy import ndimage
from scipy.ndimage import gaussian_filter
from skimage.filters import threshold_otsu
from sklearn.cluster import KMeans
from skimage import feature, color
from PIL import Image, ImageOps
from tqdm import tqdm

EXTS = ['.png', '.jpg', '.jpeg', '.bmp']


class xDog:
    def __init__(self,
                 γ: float = 0.95,
                 ϕ: float = 1e9,
                 ϵ: float = -1e1,
                 k: float = 4.5,
                 σ: float = 0.3, ) -> None:
        self.γ = γ
        self.ϕ = ϕ
        self.ϵ = ϵ
        self.k = k
        self.σ = σ

    def __call__(self, img) -> np.ndarray:
        x = (img[..., 0] + img[..., 1] + img[..., 2]) / 3

        gaussian_a = gaussian_filter(x, self.σ)
        gaussian_b = gaussian_filter(x, self.σ * self.k)

        dog = gaussian_a - self.γ * gaussian_b

        inf = dog < self.ε
        xdog = inf * 1 + ~inf * (1 - np.tanh(self.φ * dog))

        xdog -= xdog.min()
        xdog /= xdog.max()
        xdog = xdog >= threshold_otsu(xdog)
        xdog = 1 - xdog

        return xdog


class celebMask:
    def __init__(self):
        super(celebMask, self).__init__()
        self.inner_parts = ['skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'l_ear', 'r_ear', 'nose', 'u_lip', 'mouth', 'l_lip', 'eye_g', 'hair','neck']

    def imgpath_to_labelpath(self, path):
        root, name = path.replace('\\', '/').split('CelebA-HQ-img-subset15000/')
        subfolder = str(int(name.split('.')[0]) // 2000)
        label_path = os.path.join(root, 'CelebAMask-HQ-mask-anno', subfolder, name.split('.')[0].zfill(5) + '_{}.png')
        # part_label = []
        # for part in self.inner_parts:
        #     if os.path.exists(label_path.format(part)):
        #         part_label.append(label_path.format(part))
        return label_path

    def labelpath_to_imgpath(self, path):
        root= path.replace('\\', '/').split('CelebAMask-HQ-mask-anno/')[0]
        name = os.path.basename(path).split('_')[0]
        img_path = os.path.join(root, 'CelebA-HQ-img', str(int(name)) + '.jpg')
        return img_path

    def get_edges(self, edge, t):
        edge[:,1:] = edge[:,1:] | (t[:,1:] != t[:,:-1])
        edge[:,:-1] = edge[:,:-1] | (t[:,1:] != t[:,:-1])
        edge[1:,:] = edge[1:,:] | (t[1:,:] != t[:-1,:])
        edge[:-1,:] = edge[:-1,:] | (t[1:,:] != t[:-1,:])
        return edge

    def __call__(self, path):
        label_paths = self.imgpath_to_labelpath(path)
        size = (512,512)
        import matplotlib.pyplot as plt
        inner_label = np.ones(size, dtype=np.uint8)
        edges = np.zeros(size, dtype=np.uint8)
        # tensors_dist = 0
        # e = 1
        for part in self.inner_parts:
            edge = np.zeros(size, dtype=np.uint8)
            if os.path.exists(label_paths.format(part)):
                part_label = Image.open(label_paths.format(part)).convert('L').resize(size, resample=Image.NEAREST)
                part_label = np.array(part_label)
                if part == 'hair':
                    inner_label[part_label == 255] = 1
                else:
                    inner_label[part_label == 255] = 0
                edges = self.get_edges(edges, part_label)
            #     edge = self.get_edges(edge, part_label)
            # im_dist = cv2.distanceTransform(255 - edge * 255, cv2.DIST_L1, 3)
            # im_dist = np.clip((im_dist / 3), 0, 255).astype(np.uint8)
            # tensor_dist = transform_img(Image.fromarray(im_dist))
            # tensors_dist = tensor_dist if e == 1 else torch.cat([tensors_dist, tensor_dist])
            # e += 1
        # img = Image.open(path).resize(size, resample=Image.BILINEAR)
        # canny_edges = feature.canny(np.array(img.convert('L')))
        # canny_edges = canny_edges * inner_label

        # edges_all = edges + canny_edges
        # edges_all[edges_all > 1] = 1
        edges = np.ones_like(edges)*1 - edges
        return Image.fromarray((edges * 255).astype('uint8'), 'L')


def allowed_extension(path):
    target = os.path.splitext(path)[1]
    return any(ext == target for ext in EXTS)


def open_img(path):
    """
	Typically an icon image contains Alpha channel
	We paste the icon to a white background image to eliminate the alpha channel
	if we directly discard alpha channel there'll be some artifact.
	"""
    img = Image.open(path).convert('RGBA')
    assert img.size[0] == img.size[1], "Cannot handle rectangular icon image"
    nimg = Image.new(mode='RGBA', size=img.size, color=(255, 255, 255, 255))
    nimg.paste(img, mask=img.split()[3])
    return nimg.convert('RGB')


def get_Lab_3dhist(img):
    """
	return 3d Lab histogram
	shape = (bins, bins, bins)
	"""
    sample = np.array(img)
    non_white_pixels = color.rgb2lab(sample)[np.any(sample != [255, 255, 255], axis=2), :]
    # L: 0 to 100, a: -127 to 128, b: -128 to 127.
    H, _ = np.histogramdd(non_white_pixels, bins=(8, 8, 8), range=[(0, 100), (-127, 128), (-128, 127)])
    return H


def get_resized_and_contour(img, sketch_generator, args):
    """
	return the contour and icon in resized target resolution
	(image padding included)
	"""
    contour = sketch_generator(np.array(img) / 255.0)
    contour = Image.fromarray((contour * 255).astype('uint8'), 'L')
    resolution = args.resolution
    img = img.resize((resolution, resolution), Image.BICUBIC)

    contour = contour.resize((resolution, resolution), Image.BICUBIC)
    return img, contour

def getlabelpt(args, histograms):
    # kmeans
    print('Kmeans n_clusters =', args.n_clusters)
    kmeans = KMeans(n_clusters=args.n_clusters, random_state=0).fit(histograms)
    labels = kmeans.labels_
    groups = [[] for _ in range(args.n_clusters)]
    for idx, label in enumerate(labels):
        groups[label].append(idx)
    torch.save({
        'labels': labels,
        'groups': groups
    }, os.path.join(args.save_path, 'labels.pt'))

def boot_celebmask(args):
    files = [os.path.join(args.dataset, f) for f in sorted(filter(allowed_extension, os.listdir(args.dataset)))]
    img_root = os.path.join(args.save_path, 'img')
    contour_root = os.path.join(args.save_path, 'contour')
    if os.path.exists(args.save_path):
        shutil.rmtree(args.save_path)
    os.mkdir(args.save_path)
    os.mkdir(img_root)
    os.mkdir(contour_root)

    celeb = celebMask()
    histograms = []
    import matplotlib.pyplot as plt
    for idx, f in enumerate(tqdm(files, ascii=True, ncols=100)):
        label = celeb(f)
        img = open_img(f)

        # calc histogram
        H = get_Lab_3dhist(img)
        H = H.flatten() / H.sum()
        histograms.append(H)

        # save processed image
        resolution = args.resolution
        img = img.resize((resolution, resolution), Image.BICUBIC)
        label = label.resize((resolution, resolution), Image.NEAREST)
        img_path = os.path.join(img_root, '{:06d}.png'.format(idx))
        contour_path = os.path.join(contour_root, '{:06d}.png'.format(idx))
        img.save(img_path)
        label.save(contour_path)
    histograms = np.stack(histograms)
    return histograms

def boot(args):

    files = [os.path.join(args.dataset, f) for f in sorted(filter(allowed_extension, os.listdir(args.dataset)))]
    img_root = os.path.join(args.save_path, 'img')
    contour_root = os.path.join(args.save_path, 'contour')
    if os.path.exists(args.save_path):
        shutil.rmtree(args.save_path)
    os.mkdir(args.save_path)
    os.mkdir(img_root)
    os.mkdir(contour_root)

    xdog = xDog()
    histograms = []
    for idx, f in enumerate(tqdm(files, ascii=True, ncols=100)):
        img = open_img(f)

        # calc histogram
        H = get_Lab_3dhist(img)
        H = H.flatten() / H.sum()
        histograms.append(H)

        # save processed image
        resized_img, contour = get_resized_and_contour(img, xdog, args)

        img_path = os.path.join(img_root, '{:06d}.png'.format(idx))
        contour_path = os.path.join(contour_root, '{:06d}.png'.format(idx))
        resized_img.save(img_path)
        contour.save(contour_path)
    histograms = np.stack(histograms)
    return histograms

def defineArgs():
    parser = ArgumentParser()
    parser.add_argument('--dataset', type=str, default=r'D:\zhw\data\pix2pix_edge2shot\train\image')
    # parser.add_argument('--dataset', type=str,
    #                     default=r'E:\Valdataset\ref')
    parser.add_argument('--save_path', type=str, default='preprocessed_data_pix2pix')
    parser.add_argument('--n_clusters', type=int, default=500)
    parser.add_argument('--resolution', type=int, default=256)
    parser.add_argument('--pad_ratio', type=int, default=0)
    parser.add_argument('--contour_k', type=int, default=5)
    parser.add_argument('--bins', type=int, default=8)
    parser.add_argument('--canny_sigma', type=float, default=1)
    args = parser.parse_args()
    print(vars(args))
    return args

if __name__ == '__main__':
    args = defineArgs()
    h = boot(args)
    getlabelpt(args, h)
    print('done.')
